'''
    Controls:
        Move:
            Forward - W
            Backwards - S

        Strafe:
            Up - up arrow
            Down - down arrow
            Left - A
            Right - D

        Rotate:
            Left - Q
            Right - E

        Zoom:
            In - X
            Out - Z
'''
import json

import numpy as np
from PySide6.QtWidgets import QMainWindow
from pyrr import Matrix44, Quaternion, Vector3, vector
from PySide6.QtCore import Qt
from PySide6 import QtCore, QtGui, QtWidgets

import moderngl
import custom_moderngl_window as mglw
import os
import ast

vertex_shader = """
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec3 aColor;

uniform mat4 Mvp;

out vec3 ourColor;

void main()
{
	gl_Position = Mvp * vec4(aPos, 1.0f);
	ourColor = vec3(aColor);
}
"""

fragment_shader = """
#version 330 core
out vec4 FragColor;
in vec3 ourColor;

void main()
{
	FragColor = vec4(ourColor, 1.0f);
}
"""


class Camera:
    def __init__(self, ratio):
        self._rotate_roll = 0.5
        self._zoom_step = 0.5  # 视野放大/缩小灵敏度
        self._move_vertically = 0.01  # 上下移动灵敏度
        self._move_horizontally = 0.01  # 左右移动灵敏度
        self._rotate_horizontally = 0.5  # 水平旋转灵敏度
        self._rotate_vertically = 0.5  # 垂直旋转灵敏度
        self._max_ratio = 75  # 限制旋转角度

        self._z_near = 0.1
        self._z_far = 100
        self._ratio = ratio
        self._target_position = Vector3([0.0, 0.0, 1.0])

        self.reset()

    def zoom_in(self):
        self._field_of_view_degrees = self._field_of_view_degrees - self._zoom_step
        self.build_projection()

    def zoom_out(self):
        self._field_of_view_degrees = self._field_of_view_degrees + self._zoom_step
        self.build_projection()

    def move_forward(self):
        self._camera_position = self._camera_position + self._camera_front * self._move_horizontally
        self.build_look_at()

    def move_backwards(self):
        self._camera_position = self._camera_position - self._camera_front * self._move_horizontally
        self.build_look_at()

    def strafe_left(self):
        self._camera_position = self._camera_position - vector.normalize(
            self._camera_front ^ self._camera_up) * self._move_horizontally
        self.build_look_at()

    def strafe_right(self):
        self._camera_position = self._camera_position + vector.normalize(
            self._camera_front ^ self._camera_up) * self._move_horizontally
        self.build_look_at()

    def strafe_up(self):
        self._camera_position = self._camera_position + self._camera_up * self._move_vertically
        self.build_look_at()

    def strafe_down(self):
        self._camera_position = self._camera_position - self._camera_up * self._move_vertically
        self.build_look_at()

    def rotate_left(self):
        rotation = Quaternion.from_y_rotation(2 * float(self._rotate_horizontally) * np.pi / 180)
        vector_camera = -rotation * self._camera_position - self._target_position
        # 限制旋转角度
        if abs(np.arctan(vector_camera[0] / vector_camera[2]) * 180 / np.pi) > self._max_ratio:
            return
        self._camera_position = -rotation * self._camera_position
        self._camera_front = self._cameras_target - self._camera_position
        self.build_look_at()

    def rotate_right(self):
        rotation = Quaternion.from_y_rotation(-2 * float(self._rotate_horizontally) * np.pi / 180)
        vector_camera = -rotation * self._camera_position - self._target_position
        # 限制旋转角度
        if abs(np.arctan(vector_camera[0] / vector_camera[2]) * 180 / np.pi) > self._max_ratio:
            return
        self._camera_position = -rotation * self._camera_position
        self._camera_front = self._cameras_target - self._camera_position
        self.build_look_at()

    def rotate_up(self):
        rotation = Quaternion.from_x_rotation(2 * float(self._rotate_vertically) * np.pi / 180)
        vector_camera = -rotation * self._camera_position - self._target_position
        # 限制旋转角度
        if abs(np.arctan(vector_camera[1] / vector_camera[2]) * 180 / np.pi) > self._max_ratio:
            return
        self._camera_position = -rotation * self._camera_position
        self._camera_front = self._cameras_target - self._camera_position
        print(self._camera_position)
        self.build_look_at()

    def rotate_down(self):
        rotation = Quaternion.from_x_rotation(-2 * float(self._rotate_vertically) * np.pi / 180)
        vector_camera = -rotation * self._camera_position - self._target_position
        # 限制旋转角度
        if abs(np.arctan(vector_camera[1] / vector_camera[2]) * 180 / np.pi) > self._max_ratio:
            return
        self._camera_position = -rotation * self._camera_position
        self._camera_front = self._cameras_target - self._camera_position
        print(self._camera_position)
        self.build_look_at()

    def rotate_roll_left(self):
        rotation = Quaternion.from_z_rotation(2 * float(self._rotate_roll) * np.pi / 180)
        self._camera_up = rotation * self._camera_up
        self._camera_front = rotation * self._camera_front
        self.build_look_at()

    def rotate_roll_right(self):
        rotation = Quaternion.from_z_rotation(-2 * float(self._rotate_roll) * np.pi / 180)
        self._camera_up = rotation * self._camera_up
        self._camera_front = rotation * self._camera_front
        self.build_look_at()

    def reset(self):
        self._field_of_view_degrees = 20.0
        self.build_projection()

        self._camera_position = Vector3([0.0, 0.0, 5.0])
        self._camera_front = Vector3([0.0, 0.0, -4.0])
        self._camera_up = Vector3([0.0, 1.0, 0.0])
        self._cameras_target = (self._camera_position + self._camera_front)
        self.build_look_at()

    def build_look_at(self):
        self._cameras_target = (self._camera_position + self._camera_front)
        self.mat_lookat = Matrix44.look_at(
            self._camera_position,
            self._cameras_target,
            self._camera_up)

    def build_projection(self):
        self.mat_projection = Matrix44.perspective_projection(
            self._field_of_view_degrees,
            self._ratio,
            self._z_near,
            self._z_far)


class CSVDataParse:
    def __init__(self, csv_2d, csv_3d):
        self.csv_2d = csv_2d
        self.csv_3d = csv_3d

    @staticmethod
    def readCSVData(filename):
        data = []
        try:
            with open(filename, 'r') as f:
                lines = f.readlines()
            for line in lines:
                row = line.strip().split(',')
                data.append([float(x) for x in row])
        except Exception as e:
            print(e)
        finally:
            return np.array(data)

    @staticmethod
    def normalize_with_custom_range(arr, custom_min, custom_max):
        # 计算数组的最小值和最大值
        arr_min = np.min(arr)
        arr_max = np.max(arr)
        # 对数组进行归一化到指定范围内
        normalized_arr = (arr - arr_min) * (custom_max - custom_min) / (arr_max - arr_min) + custom_min
        return normalized_arr

    def parse_data(self):
        points3D_array = np.rot90(self.readCSVData(self.csv_3d), 3)
        points2D_array = np.rot90(self.readCSVData(self.csv_2d), 3)
        pointsColor = []
        points = []
        validPoints = []
        invalidPoints = []
        indices = []
        # 不能用numpy，循环numpy.append性能太低
        for row in range(points2D_array.shape[0]):
            for col in range(points2D_array.shape[1]):
                if row < points2D_array.shape[0] - 1 and col < points2D_array.shape[1] - 1:
                    indices.append(row * points2D_array.shape[1] + col)
                    indices.append(row * points2D_array.shape[1] + col + 1)
                    indices.append((row + 1) * points2D_array.shape[1] + col)
                    indices.append((row + 1) * points2D_array.shape[1] + col + 1)
                    indices.append(row * points2D_array.shape[1] + col + 1)
                    indices.append((row + 1) * points2D_array.shape[1] + col)

                # 分离有效/无效的数据
                if points2D_array[row][col] != 255:
                    # point = glm.vec3(row, col, points3D_array[row][col])
                    validPoints.append([row, col, points3D_array[row][col]])
                    pointsColor.append([0, 0, 0])
                else:
                    invalidPoints.append([row, col, points3D_array[row][col]])
                    pointsColor.append([255, 255, 255])

        validPoints = np.array(validPoints)
        invalidPoints = np.array(invalidPoints)
        # validPointsCoff = self.normalize_with_custom_range(validPoints[:, [2]], 0.995, 1)
        # validPoints[:, 2] = validPoints[:, 2] * validPointsCoff[:, 0]

        for validPoint in validPoints:
            points3D_array[int(validPoint[0]), int(validPoint[1])] = validPoint[2]

        indexmin = 0
        for row in range(points3D_array.shape[0]):
            for col in range(points3D_array.shape[1]):
                points.append([row, col, points3D_array[row][col]])
                if row == validPoints[0][0] and col == validPoints[0][1]:
                    print(len(points))
                    indexmin = len(points) - 1

        points = np.array(points, dtype='f4')
        points[:, 0] = self.normalize_with_custom_range(points[:, 0], -0.5, 0.5)
        points[:, 1] = self.normalize_with_custom_range(points[:, 1], -0.5, 0.5)
        points[-1, 2] = 0
        points[:, 2] = self.normalize_with_custom_range(points[:, 2], 0, 1)
        points[-1, 2] = points[-2, 2]

        validPoints[:, [2]] = self.normalize_with_custom_range(validPoints[:, [2]], 200, 0)
        invalidPoints[:, [2]] = self.normalize_with_custom_range(invalidPoints[:, [2]], 250, 255)
        indexValid = 0
        indexInvalid = 0
        for index, value in enumerate(pointsColor):
            if value[0] == 0:
                pointsColor[index] = [validPoints[indexValid][2], validPoints[indexValid][2],
                                      validPoints[indexValid][2]]
                indexValid += 1
            else:
                pointsColor[index] = [invalidPoints[indexInvalid][2], invalidPoints[indexInvalid][2],
                                      invalidPoints[indexInvalid][2]]
                indexInvalid += 1
        pointsColor = np.array(pointsColor, dtype='f4')
        pointsColor = self.normalize_with_custom_range(pointsColor, 0, 1)

        indices = np.array(indices, dtype='i4')
        return points, pointsColor, indices


class PerspectiveProjection(mglw.WindowConfig):
    window_width = 500
    window_height = 500
    gl_version = (3, 3)
    title = "ModernGL"
    window_size = (window_width, window_height)
    aspect_ratio = window_width / window_height
    resizable = True

    def __init__(self, **kwargs):
        print('PerspectiveProjection init', kwargs)
        super().__init__(**kwargs)

        self.prog = self.ctx.program(
            vertex_shader=vertex_shader,
            fragment_shader=fragment_shader
        )

        self.camera = Camera(self.aspect_ratio)

        # if 'extra_args' in kwargs and '2DDataPath' in kwargs['extra_args'] and '3DDataPath' in kwargs['extra_args']:
        try:
            json_str = kwargs['extra_args'].replace("'", "\"")
            filePaths = json.loads(json_str)
            self.csvParser = CSVDataParse(filePaths['2DDataPath'], filePaths['3DDataPath'])
        except Exception as e:
            print('Arguments error:', e)
            self.csvParser = CSVDataParse(os.path.join(os.path.curdir, 'data', '0424_2D.csv'),
                                          os.path.join(os.path.curdir, 'data', '0424_3D.csv'))
        points, self.pointsColor, indices = self.csvParser.parse_data()

        self.mvp = self.prog['Mvp']
        self.vbo_pos = self.ctx.buffer(points.tobytes())
        self.vbo_color = self.ctx.buffer(self.pointsColor.tobytes())
        ebo = self.ctx.buffer(indices.tobytes())

        self.vao = self.ctx.vertex_array(
            program=self.prog,
            content=[
                (self.vbo_pos, "3f", "aPos"),
                (self.vbo_color, "3f", "aColor"),
            ],
            index_buffer=ebo,
            index_element_size=4
        )

        self.states = {
            self.wnd.keys.W: False,  # forward
            self.wnd.keys.S: False,  # backwards
            self.wnd.keys.UP: False,  # strafe Up
            self.wnd.keys.DOWN: False,  # strafe Down
            self.wnd.keys.A: False,  # strafe left
            self.wnd.keys.D: False,  # strafe right
            self.wnd.keys.Q: False,  # rotate left
            self.wnd.keys.E: False,  # rotare right
            self.wnd.keys.Z: False,  # zoom in
            self.wnd.keys.X: False,  # zoom out
        }
        self.mouse_state = Qt.MouseButton.NoButton

    def move_camera(self):
        if self.states.get(self.wnd.keys.W):
            self.camera.move_forward()

        if self.states.get(self.wnd.keys.S):
            self.camera.move_backwards()

        if self.states.get(self.wnd.keys.UP):
            self.camera.strafe_up()

        if self.states.get(self.wnd.keys.DOWN):
            self.camera.strafe_down()

        if self.states.get(self.wnd.keys.A):
            self.camera.strafe_left()

        if self.states.get(self.wnd.keys.D):
            self.camera.strafe_right()

        if self.states.get(self.wnd.keys.Q):
            self.camera.rotate_roll_left()

        if self.states.get(self.wnd.keys.E):
            self.camera.rotate_roll_right()

        if self.states.get(self.wnd.keys.Z):
            self.camera.zoom_in()

        if self.states.get(self.wnd.keys.X):
            self.camera.zoom_out()

    def key_event(self, key, action, modifiers):
        if key not in self.states:
            print(key, action)
            return

        if action == self.wnd.keys.ACTION_PRESS:
            self.states[key] = True
        else:
            self.states[key] = False

    def mouse_drag_event(self, x: int, y: int, dx: int, dy: int):
        print(self.mouse_state, x, y, dx, dy)
        if self.mouse_state == 1:
            if dx > 0:
                self.camera.strafe_left()
            elif dx < 0:
                self.camera.strafe_right()
            if dy > 0:
                self.camera.strafe_up()
            elif dy < 0:
                self.camera.strafe_down()
        elif self.mouse_state == 2:
            if dx > 0:
                self.camera.rotate_right()
            elif dx < 0:
                self.camera.rotate_left()
            if dy > 0:
                self.camera.rotate_down()
            elif dy < 0:
                self.camera.rotate_up()

    def mouse_press_event(self, x: int, y: int, button: int):
        if button == 3:
            self.camera.reset()
        else:
            self.mouse_state = button

    def mouse_release_event(self, x: int, y: int, button: int):
        self.mouse_state = Qt.MouseButton.NoButton

    def mouse_scroll_event(self, x_offset: float, y_offset: float):
        if y_offset > 0:
            self.camera.zoom_in()
        if y_offset < 0:
            self.camera.zoom_out()

    # rgb to gray
    def switch_color_to_gray_event(self):
        print('switch_color_to_gray_event')

    # gray to rgb
    def switch_color_to_rgb_event(self):
        print('switch_color_to_rgb_event')
        # 标记需要保持不变的点
        keep_same = np.all(self.pointsColor > 0.98, axis=1)
        filtered_pointsColor = self.pointsColor[~keep_same]
        gray_ch_color = filtered_pointsColor[:, 0]
        gray_ch_color = CSVDataParse.normalize_with_custom_range(gray_ch_color, -1.0, 1.0)
        gray_index = 0
        for index, color in enumerate(self.pointsColor):
            if not keep_same[index]:
                if color[0] != color[1] or color[1] != color[2]:
                    return
                if gray_ch_color[gray_index] > 0:
                    self.pointsColor[index] = [gray_ch_color[gray_index], 0.0, 0.0]
                else:
                    self.pointsColor[index] = [0.0, 0.0, -gray_ch_color[gray_index]]
                gray_index = gray_index + 1

    def render(self, time, frame_time):
        self.move_camera()

        self.ctx.clear(.0, .0, .4)
        self.ctx.enable(moderngl.DEPTH_TEST)

        self.vbo_color.write(self.pointsColor.tobytes())
        self.mvp.write((self.camera.mat_projection * self.camera.mat_lookat).astype('f4'))
        self.vao.render(moderngl.TRIANGLES)


if __name__ == '__main__':
    app = QtWidgets.QApplication([])
    PerspectiveProjection.run()
